package com.tool4j.service;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.TranslationOptions;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.fastjson.JSON;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.tool4j.common.ModelEnum;
import com.tool4j.common.ToolTypeEnum;
import com.tool4j.entity.AiSessionLog;
import com.tool4j.entity.TermInfo;
import com.tool4j.entity.params.ChatParams;
import com.tool4j.entity.params.PageParams;
import com.tool4j.entity.params.ToolCommonParams;
import com.tool4j.entity.vo.DMessage;
import com.tool4j.graph.GenerateCodeGraph;
import com.tool4j.sse.SseEventPublisher;
import com.tool4j.tools.ConversionTool;
import com.tool4j.util.UserPermission;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Component
@Slf4j
public class ChatService {

    @Autowired
    private SseEventPublisher eventPublisher;

    private final ChatClient chatClient;

    private final ChatClient chatMemoryClient;

    @Autowired
    private GenerateCodeGraph generateCodeGraph;

    @Autowired
    private AiSessionLogService sessionLogService;

    @Autowired
    private AiChatLogService chatLogService;

    @Autowired
    private TermInfoService termInfoService;

    private final Map<String, ChatClient> chatClientMap = new HashMap<>();

    /**
     * 构造方法，构建ChatClient
     *
     * @param builder
     */
    public ChatService(ChatClient.Builder builder) {
        // 默认对话客户端
        this.chatClient = builder.build();
        //带记忆的对话客户端
        ChatMemory chatMemory = MessageWindowChatMemory.builder()
                .maxMessages(10)
                .build();
        this.chatMemoryClient = builder.defaultAdvisors(
                MessageChatMemoryAdvisor.builder(chatMemory).build()
        ).build();
        // 加载各种模型
        loadModels(chatMemory);
    }

    /**
     * 加载各种模型
     *
     * @param chatMemory
     */
    public void loadModels(ChatMemory chatMemory) {
        for (ModelEnum modelEnum : ModelEnum.values()) {
            OpenAiApi openAiApi;
            if (modelEnum.getModelName().equals(ModelEnum.ERNIE_BOT.getModelName())
                    || modelEnum.getModelName().equals(ModelEnum.ERNIE_BOT_PRO.getModelName())
            ) {
                openAiApi = OpenAiApi.builder()
                        .baseUrl(modelEnum.getBaseUrl())
                        .apiKey(modelEnum.getApiKey())
                        .completionsPath("/v2/chat/completions")
                        .build();
            } else if (modelEnum.getModelName().equals(ModelEnum.CHAT_GLM.getModelName())) {
                openAiApi = OpenAiApi.builder()
                        .baseUrl(modelEnum.getBaseUrl())
                        .apiKey(modelEnum.getApiKey())
                        .completionsPath("/v4/chat/completions")
                        .build();
            } else {
                openAiApi = OpenAiApi.builder()
                        .baseUrl(modelEnum.getBaseUrl())
                        .apiKey(modelEnum.getApiKey())
                        .build();
            }
            OpenAiChatModel chatModel = OpenAiChatModel.builder().openAiApi(openAiApi).defaultOptions(
                    OpenAiChatOptions.builder().model(modelEnum.getModelName()).build()
            ).build();
            ChatClient client = ChatClient.builder(chatModel)
                    .defaultAdvisors(MessageChatMemoryAdvisor.builder(chatMemory).build())
                    .build();
            chatClientMap.put(modelEnum.getModelName(), client);
        }
    }

    /**
     * AI对话
     *
     * @param params
     */
    public void chatStream(ChatParams params) {
        String toolType = params.getToolType();
        ToolTypeEnum typeEnum = ToolTypeEnum.getByCode(toolType);
        if (toolType == null || typeEnum == null) {
            //会话如果不存在就新增
            sessionLogService.saveIfNotExists(params);
            //新增对话
            chatLogService.saveLog(params);
            normalChat(params);
            return;
        }
        chat(params, typeEnum);
    }

    private void chat(ChatParams params, ToolTypeEnum typeEnum) {
        switch (typeEnum) {
            case CRUD:
                generateCodeGraph.generate(params);
                break;
            case CODE:
                codeChat(params);
                break;
            case TRANSLATION:
                translationChat(params);
                break;
            case CONVERSION:
                conversion(params);
                break;
            default:
                normalChat(params);
        }
    }

    /**
     * 正常对话
     *
     * @param params
     */
    private void normalChat(ChatParams params) {
        String content = params.getContent();
        if (StrUtil.isNotBlank(params.getFileContent())) {
            content = String.format("%s \n\n 附带文件内容：\n%s", content, params.getFileContent());
        }
        String chatId = params.getChatId();
        String clientId = params.getClientId();
        String sessionId = params.getSessionId();
        ChatClient client = chatClientMap.get(params.getModelName());
        if (client == null) {
            client = this.chatMemoryClient;
        }

        client.prompt()
                .user(content)
                .advisors(spec -> spec.param(
                        ChatMemory.CONVERSATION_ID, sessionId))
                .stream()
                .chatResponse()
                .doOnComplete(() -> {
                    // 流正常结束后，发送结束通知
                    eventPublisher.sendComplete(params.getClientId());
                })
                .subscribe(chatResponse -> {
                    try {
                        String result = chatResponse.getResult().getOutput().getText();
                        if (result != null) {
                            String resultMessage = JSON.toJSONString(Map.of("chatId", chatId, "content", result));
                            chatLogService.updateLog(chatId, result);
                            eventPublisher.sendMessageToClient(clientId, resultMessage);
                        }
                    } catch (Exception e) {
                        log.error("输出消息异常", e);
                    }
                });
    }

    private void codeChat(ChatParams params) {
        String content = params.getContent();
        if (StrUtil.isNotBlank(params.getFileContent())) {
            content = String.format("%s \n\n 附带文件内容：\n%s", content, params.getFileContent());
        }
        String chatId = params.getChatId();
        String clientId = params.getClientId();
        String sessionId = params.getSessionId();
        String prompt =
                String.format("你是一个擅长%s编程的专业编程工具，请根据用户需求，深入思考，然后专业、精准的回答用户的问题"
                        , params.getToolParams().get("codeType"));
        chatClientMap.get(ModelEnum.QWEN_CODER.getModelName())
                .prompt(prompt)
                .user(content)
                .advisors(spec -> spec.param(
                        ChatMemory.CONVERSATION_ID, sessionId))
                .stream()
                .chatResponse()
                .doOnComplete(() -> {
                    // 流正常结束后，发送结束通知
                    eventPublisher.sendComplete(params.getClientId());
                })
                .subscribe(chatResponse -> {
                    String result = chatResponse.getResult().getOutput().getText();
                    String resultMessage = JSON.toJSONString(Map.of("chatId", chatId, "content", result));
                    eventPublisher.sendMessageToClient(clientId, resultMessage);
                });
    }

    /**
     * 翻译
     *
     * @param params
     */
    private void translationChat(ChatParams params) {
        TranslationOptions.TranslationOptionsBuilder<?, ?> builder = TranslationOptions.builder()
                .sourceLang(params.getToolParams().get("translationSource").toString())
                .targetLang(params.getToolParams().get("translationTarget").toString());
        /**
         * 术语表
         */
        if (params.getToolParams().get("terms") != null) {
            List<TranslationOptions.Term> terms = (List<TranslationOptions.Term>) params.getToolParams().get("terms");
            builder.terms(terms);
        }
        TranslationOptions options = builder.build();

        List<Message> messages = new ArrayList<>();
        String content = params.getContent();
        if (StrUtil.isNotBlank(params.getFileContent())) {
            content = String.format("%s \n\n 附带文件内容：\n%s", content, params.getFileContent());
        }
        messages.add(Message.builder().role(Role.USER.getValue()).content(content).build());
        GenerationParam param = GenerationParam.builder()
                // 若没有配置环境变量，请用阿里云百炼API Key将下行替换为：.apiKey("sk-xxx")
                .apiKey(ModelEnum.QWEN_MT_TURBO.getApiKey())
                .model(ModelEnum.QWEN_MT_TURBO.getModelName())
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .translationOptions(options)
                .incrementalOutput(true)
                .messages(messages)
                .build();
        Generation gen = new Generation();
        String chatId = params.getChatId();
        String clientId = params.getClientId();
        try {
            gen.streamCall(param)
                    .doOnComplete(() -> {
                        // 流正常结束后，发送结束通知
                        eventPublisher.sendComplete(params.getClientId());
                    })
                    .subscribe(message -> {
                        String result = message.getOutput().getChoices().get(0).getMessage().getContent();
                        String resultMessage = JSON.toJSONString(Map.of("chatId", chatId, "content", result));
                        eventPublisher.sendMessageToClient(clientId, resultMessage);
                    });
        } catch (Exception e) {
            log.error("翻译异常", e);
        }
    }

    /**
     * 格式转换
     *
     * @param params
     */
    private void conversion(ChatParams params) {
        String content = params.getContent();
        if (StrUtil.isNotBlank(params.getFileContent())) {
            content = String.format("%s \n\n 附带文件内容：\n%s", content, params.getFileContent());
        }
        String chatId = params.getChatId();
        String clientId = params.getClientId();
        String sessionId = params.getSessionId();
        String prompt =
                String.format("你是一个专业的格式转换工具，请将用户的原始文本（%s）转换为：%s，输出转换后的内容，不要有额外的说明"
                        , params.getToolParams().get("conversionSource")
                        , params.getToolParams().get("conversionTarget"));
        this.chatClient
                .prompt(prompt)
                .user(content)
                .advisors(spec -> spec.param(
                        ChatMemory.CONVERSATION_ID, sessionId))
                .tools(new ConversionTool())
                .stream()
                .chatResponse()
                .doOnComplete(() -> {
                    // 流正常结束后，发送结束通知
                    eventPublisher.sendComplete(params.getClientId());
                })
                .subscribe(chatResponse -> {
                    String result = chatResponse.getResult().getOutput().getText();
                    String resultMessage = JSON.toJSONString(Map.of("chatId", chatId, "content", result));
                    eventPublisher.sendMessageToClient(clientId, resultMessage);
                });
    }

    public Page<AiSessionLog> getSessions(ServerWebExchange exchange, PageParams<AiSessionLog> params) {
        Long userId = UserPermission.get(exchange).getUserId();
        return sessionLogService.findPage(userId, params);
    }

    public void deleteSession(ServerWebExchange exchange, Long id) {
        Long userId = UserPermission.get(exchange).getUserId();
        sessionLogService.deleteSession(userId, id);
    }

    public void deleteAllSession(ServerWebExchange exchange) {
        Long userId = UserPermission.get(exchange).getUserId();
        sessionLogService.deleteAllSession(userId);
    }

    public List<DMessage> getMessages(Long sessionId) {
        return chatLogService.findList(sessionId);
    }

    public void translation(Long userId, ToolCommonParams params) {
        TranslationOptions.TranslationOptionsBuilder<?, ?> builder = TranslationOptions.builder()
                .sourceLang(params.getSource())
                .targetLang(params.getTarget());
        /**
         * 术语表
         */
        List<TermInfo> termInfos = termInfoService.findAll(userId);
        if (CollUtil.isNotEmpty(termInfos)) {
            List<TranslationOptions.Term> terms = new ArrayList<>();
            for (TermInfo termInfo : termInfos) {
                TranslationOptions.Term term = TranslationOptions.Term.builder()
                        .source(termInfo.getSource())
                        .target(termInfo.getTarget())
                        .build();
                terms.add(term);
            }
            builder.terms(terms);
        }
        TranslationOptions options = builder.build();

        List<Message> messages = new ArrayList<>();
        String content = params.getContent();
        messages.add(Message.builder().role(Role.USER.getValue()).content(content).build());
        GenerationParam param = GenerationParam.builder()
                // 若没有配置环境变量，请用阿里云百炼API Key将下行替换为：.apiKey("sk-xxx")
                .apiKey(ModelEnum.QWEN_MT_TURBO.getApiKey())
                .model(ModelEnum.QWEN_MT_TURBO.getModelName())
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .translationOptions(options)
                .incrementalOutput(true)
                .messages(messages)
                .build();
        Generation gen = new Generation();
        try {
            gen.streamCall(param)
                    .doOnComplete(() -> {
                        // 流正常结束后，发送结束通知
                        eventPublisher.sendComplete(params.getClientId());
                    })
                    .subscribe(message -> {
                        String result = message.getOutput().getChoices().get(0).getMessage().getContent();
                        String resultMessage = JSON.toJSONString(Map.of("chatId", params.getChatId(), "content", result));
                        eventPublisher.sendMessageToClient(params.getClientId(), resultMessage);
                    });
        } catch (Exception e) {
            log.error("翻译出错", e);
        }
    }

    public void conversion(ToolCommonParams params) {
        String content = params.getContent();
        String chatId = params.getChatId();
        String clientId = params.getClientId();
        String prompt =
                String.format("你是一个专业的格式转换工具，请将用户的原始文本（%s）转换为：%s，输出转换后的内容，不要有额外的说明"
                        , params.getSource()
                        , params.getTarget());
        this.chatClient
                .prompt(prompt)
                .user(content)
                .tools(new ConversionTool())
                .stream()
                .chatResponse()
                .doOnComplete(() -> {
                    // 流正常结束后，发送结束通知
                    eventPublisher.sendComplete(params.getClientId());
                })
                .subscribe(chatResponse -> {
                    String result = chatResponse.getResult().getOutput().getText();
                    String resultMessage = JSON.toJSONString(Map.of("chatId", chatId, "content", result));
                    eventPublisher.sendMessageToClient(clientId, resultMessage);
                });
    }

    public void generateComment(ToolCommonParams params) {
        String prompt = "你是一个专业的代码注释生成工具，请为用户的代码添加注释，不允许修改代码原始逻辑，保持原来的代码格式，输出转换后的内容，不要有额外的说明";
        String chatId = params.getChatId();
        String clientId = params.getClientId();
        chatClientMap.get(ModelEnum.QWEN_CODER.getModelName())
                .prompt(prompt)
                .user(params.getContent())
                .advisors(spec -> spec.param(
                        ChatMemory.CONVERSATION_ID, params.getSessionId()))
                .stream()
                .chatResponse()
                .doOnComplete(() -> {
                    // 流正常结束后，发送结束通知
                    eventPublisher.sendComplete(params.getClientId());
                })
                .subscribe(chatResponse -> {
                    String result = chatResponse.getResult().getOutput().getText();
                    String resultMessage = JSON.toJSONString(Map.of("chatId", chatId, "content", result));
                    eventPublisher.sendMessageToClient(clientId, resultMessage);
                });
    }
}
